#!/usr/bin/env python
# -*- coding: utf-8 -*-

import socket
import threading
import time
from typing import Optional, Callable, Dict
from logger import logger
from utils import Encryptor
from config import config

class NetworkManager:
    def __init__(self):
        self.socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        self.socket.bind(('0.0.0.0', 0))
        self.running = False
        self.listen_thread = None
        self.message_handlers: Dict[str, Callable] = {}
        self.encryptor = None
        self.remote_address = None

    def start(self) -> None:
        """启动网络管理器"""
        self.running = True
        self.listen_thread = threading.Thread(target=self._listen)
        self.listen_thread.daemon = True
        self.listen_thread.start()
        logger.info('网络管理器启动')

    def stop(self) -> None:
        """停止网络管理器"""
        self.running = False
        if self.socket:
            self.socket.close()
        if self.listen_thread:
            self.listen_thread.join()
        logger.info('网络管理器停止')

    def set_encryption_key(self, key: bytes) -> None:
        """设置加密密钥"""
        self.encryptor = Encryptor(key)

    def set_remote_address(self, address: tuple) -> None:
        """设置远程节点地址"""
        self.remote_address = address

    def register_handler(self, message_type: str, handler: Callable) -> None:
        """注册消息处理器"""
        self.message_handlers[message_type] = handler

    def send_message(self, message_type: str, data: bytes) -> bool:
        """发送消息到远程节点"""
        if not self.remote_address:
            logger.error('未设置远程节点地址')
            return False

        try:
            # 构造消息格式：类型:加密数据
            if self.encryptor:
                encrypted_data = self.encryptor.encrypt(data)
            else:
                encrypted_data = data

            message = f'{message_type}:'.encode() + encrypted_data
            self.socket.sendto(message, self.remote_address)
            return True
        except Exception as e:
            logger.error(f'发送消息失败: {e}')
            return False

    def _listen(self) -> None:
        """监听网络消息"""
        while self.running:
            try:
                data, addr = self.socket.recvfrom(65535)  # 最大UDP包大小
                if data:
                    self._handle_message(data, addr)
            except Exception as e:
                if self.running:
                    logger.error(f'接收消息失败: {e}')

    def _handle_message(self, data: bytes, addr: tuple) -> None:
        """处理接收到的消息"""
        try:
            # 解析消息类型和数据
            message_parts = data.split(b':', 1)
            if len(message_parts) != 2:
                return

            message_type = message_parts[0].decode()
            message_data = message_parts[1]

            # 解密数据
            if self.encryptor:
                message_data = self.encryptor.decrypt(message_data)

            # 调用对应的消息处理器
            handler = self.message_handlers.get(message_type)
            if handler:
                handler(message_data, addr)
            else:
                logger.warning(f'未找到消息类型 {message_type} 的处理器')

        except Exception as e:
            logger.error(f'处理消息失败: {e}')